#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import json
from typing import List

from agentscope.agent import AgentBase
from agentscope.message import Msg
from agentscope.model import DashScopeChatModel
from agentscope.formatter import DashScopeChatFormatter
from agentscope.memory import InMemoryMemory

from src.utils.env_config import aliyun_model_config

class AbstractGeneratorAgent(AgentBase):
    """T3 - 摘要生成智能体"""

    def __init__(self):
        super().__init__()
        self.name = "AbstractGeneratorAgent"
        self.sys_prompt = """你是一个摘要生成智能体。你的任务是基于文献检索结果，撰写一份正式的报告摘要。
        要求：
        1. 内容来源于T2的检索结果摘要，不得虚构未提及的信息；
        2. 摘要长度不少于100字；
        3. 语言正式、逻辑清晰，包含背景、核心发现、趋势判断；
        4. 使用第三人称叙述，适合写入正式报告；
        5. 输出格式必须为JSON，仅包含一个字段"abstract"，值为字符串；
        6. 不要解释过程，不要添加额外内容。
        示例输出：{"abstract": "本报告基于对2025年第三季度人工智能领域相关文献的系统性检索与分析，总结了当前AI技术发展的主要趋势。
        研究发现，大模型轻量化、多模态融合与边缘部署成为技术创新的核心方向……"}"""

        self.model = DashScopeChatModel(
            model_name=aliyun_model_config["model_name"],
            api_key=aliyun_model_config["api_key"],
            stream=False
        )
        self.formatter = DashScopeChatFormatter()
        self.memory = InMemoryMemory()

    async def reply(self, msg: Msg) -> Msg:
        """处理摘要生成请求"""
        await self.memory.add(msg)

        # 准备提示
        prompt = await self.formatter.format([
            Msg("system", self.sys_prompt, "system"),
            msg
        ])

        # 调用模型
        response = await self.model(prompt)

        # 创建回复消息
        reply_msg = Msg(
            name=self.name,
            content=response.content,
            role="assistant"
        )

        await self.memory.add(reply_msg)
        return reply_msg

    async def generate_abstract(self, retrieval_summary: str) -> str:
        """生成摘要"""
        user_msg = Msg(
            name="user",
            content=f"现在开始处理：\n检索结果摘要：{retrieval_summary}",
            role="user"
        )

        response_msg = await self.reply(user_msg)
        try:
            # 提取文本内容
            content_text = ""
            if isinstance(response_msg.content, list):
                for block in response_msg.content:
                    if hasattr(block, 'text'):
                        content_text += block.text
                    elif isinstance(block, dict) and 'text' in block:
                        content_text += block['text']
            else:
                content_text = str(response_msg.content)

            result = json.loads(content_text)
            return result.get("abstract", "")
        except json.JSONDecodeError:
            # 如果JSON解析失败，返回默认摘要
            print(f"JSON解析失败:{response_msg.content}")
            return ""

    async def observe(self, msg: Msg) -> None:
        """观察消息"""
        await self.memory.add(msg)

    async def handle_interrupt(self) -> Msg:
        """处理中断"""
        return Msg(
            name=self.name,
            content="摘要生成任务被中断",
            role="assistant"
        )

